// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Threading.Tasks;
using System.Threading.Tasks.Sources;

namespace Microsoft.AspNetCore.ConcurrencyLimiter;

/// <summary>
/// Custom awaiter to allow the StackPolicy to reduce allocations.
/// When this completion source has its result checked, it resets itself and returns itself to the cache of its parent StackPolicy.
/// Then when the StackPolicy needs a new completion source, it tries to get one from its cache, otherwise it allocates.
/// </summary>
internal class ResettableBooleanCompletionSource : IValueTaskSource<bool>
{
    ManualResetValueTaskSourceCore<bool> _valueTaskSource;
    private readonly StackPolicy _queue;

    public ResettableBooleanCompletionSource(StackPolicy queue)
    {
        _queue = queue;
        _valueTaskSource.RunContinuationsAsynchronously = true;
    }

    public ValueTask<bool> GetValueTask()
    {
        return new ValueTask<bool>(this, _valueTaskSource.Version);
    }

    bool IValueTaskSource<bool>.GetResult(short token)
    {
        var isValid = token == _valueTaskSource.Version;
        try
        {
            return _valueTaskSource.GetResult(token);
        }
        finally
        {
            if (isValid)
            {
                _valueTaskSource.Reset();
                _queue._cachedResettableTCS = this;
            }
        }
    }

    public ValueTaskSourceStatus GetStatus(short token)
    {
        return _valueTaskSource.GetStatus(token);
    }

    void IValueTaskSource<bool>.OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
    {
        _valueTaskSource.OnCompleted(continuation, state, token, flags);
    }

    public void Complete(bool result)
    {
        _valueTaskSource.SetResult(result);
    }
}
